package com.smmisia.auth.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.smmisia.auth.aop.Encrypted;
import com.smmisia.auth.utls.AESUtil;
import com.smmisia.common.Result.Result;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

/**
 * @author smmisia
 * @date 2024/4/16 16:44
 * @description:加密响应数据
 */
@RestControllerAdvice
public class SecretResponseAdvice implements ResponseBodyAdvice<Object> {

    private static final String SECRET_KEY = "smmisia"; // 替换为你的密钥
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        // 判断是否需要加密（通过注解判断）
        return returnType.hasMethodAnnotation(Encrypted.class);
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
                                  Class<? extends HttpMessageConverter<?>> selectedConverterType,
                                  org.springframework.http.server.ServerHttpRequest request,
                                  org.springframework.http.server.ServerHttpResponse response) {
        try {
            // 检查返回值是否是 Result 类型
            if (body instanceof Result<?>) {
                Result<?> result = (Result<?>) body;

                // 加密 data 字段
                if (result.getData() != null) {
                    String dataJson = objectMapper.writeValueAsString(result.getData());
                    String encryptedData = AESUtil.encrypt(dataJson, SECRET_KEY);
                    // 使用强制类型转换将 data 设置为加密后的 String
                    @SuppressWarnings("unchecked")
                    Result<String> encryptedResult = (Result<String>) result;
                    encryptedResult.setData(encryptedData);
                    return encryptedResult;
                }

                return result;
            }
            // 如果返回值不是 Result 类型，按普通字符串加密
            return AESUtil.encrypt(body.toString(), SECRET_KEY);
        } catch (Exception e) {
            throw new RuntimeException("Failed to encrypt response body", e);
        }
    }
}
